use axum::{
    extract::{Request, State},
    http::{HeaderMap, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
};
use std::net::SocketAddr;
use std::sync::Arc;
use tracing::{info, warn};

use crate::ip_auth::IpAuthService;

/// IP认证中间件
pub async fn ip_auth_middleware(
    State(ip_auth_service): State<Arc<IpAuthService>>,
    request: Request,
    next: Next,
) -> Response {
    // 获取客户端IP地址
    let client_ip = extract_client_ip(&request);
    
    // 检查IP访问权限
    if let Err(error) = ip_auth_service.check_ip_access(&client_ip) {
        warn!("IP access denied: {} - {}", client_ip, error);
        return (StatusCode::FORBIDDEN, format!("IP access denied: {}", error)).into_response();
    }
    
    info!("IP access allowed: {}", client_ip);
    next.run(request).await
}

/// 提取客户端IP地址
fn extract_client_ip(request: &Request) -> std::net::IpAddr {
    // 首先尝试从X-Forwarded-For头部获取（代理场景）
    if let Some(forwarded_for) = get_forwarded_for_ip(request.headers()) {
        return forwarded_for;
    }
    
    // 然后尝试从X-Real-IP头部获取
    if let Some(real_ip) = get_real_ip(request.headers()) {
        return real_ip;
    }
    
    // 最后从连接信息中获取
    if let Some(remote_addr) = request.extensions().get::<SocketAddr>() {
        return remote_addr.ip();
    }
    
    // 默认返回本地IP
    "127.0.0.1".parse().unwrap()
}

/// 从X-Forwarded-For头部提取IP
fn get_forwarded_for_ip(headers: &HeaderMap) -> Option<std::net::IpAddr> {
    headers.get("x-forwarded-for")
        .and_then(|value| value.to_str().ok())
        .and_then(|value| {
            // 取第一个IP（可能是代理链）
            value.split(',')
                .next()
                .and_then(|ip| ip.trim().parse().ok())
        })
}

/// 从X-Real-IP头部提取IP
fn get_real_ip(headers: &HeaderMap) -> Option<std::net::IpAddr> {
    headers.get("x-real-ip")
        .and_then(|value| value.to_str().ok())
        .and_then(|ip| ip.parse().ok())
}